# -*- coding: utf-8 -*-
import pexpect, re, time
from exceptions import Exception
from pexpect import ExceptionPexpect
MSFTIMEOUT = 120

class TIMEOUT(Exception):
    """
    TIMEOUT Exception objects are thrown, when metasploit is executing
    command for too long. Mainly thrown after catching pexpect.TIMEOUT. Despite
    that, it was implemented to clarify source of the errors.
    """
        
    def __init__(self):
        pass
    def __str__(self):
        return "Timeout while executing of metasploit command"
class PluginError(Exception):
    def __init__(self,value):
        self.value = value
    def __str__(self):
        return repr(self.value)
class ConnectionError(Exception):
    def __init__(self):
        pass
    def __str__(self):
        return "Not connected to Metasploit console"
class MsfExecNotFound(Exception):
    def __init__(self):
        pass
    def __str__(self):
        return "MetaSploit Executable not found (possible incorrect path)"

class MetaSploit:
    """
    Base class containing methods responsible for communication with Metasploit Framework
    """
    def __init__(self):
        """
        Default constructor for the Metasploit base class is (abstract) empty.
        """
        pass 
    def __del__(self):
        """
        Destructor for MetaSploit class. Kills the child process spawned
        by our program. In most cases it is a telnet process or localy running
        MetaSploit Framework""" 
        try:
            self.pipe.kill(2)
        except AttributeError:
            pass
    def showMatchingExploits(self):
        """
        Performs matching of available MetaSploit exploits with relation to
        services information in the database. Matching is based on IP 
        destination ports of a particular services and information
        provided by authors of MetaSploit exploits. 
        """
        if self.connected != True:
            raise ConnectionError
        try:
            self.pipe.sendline("db_autopwn -t -p")
            self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
            raise TIMEOUT
        lines =  self.pipe.before.split("\r\n")
        #return "\n".join(lines[2:])
        ret = []
        for line in lines:
            rexpl = re.match("^\[\*\] Matched (.*) against (\d+\.\d+\.\d+\.\d+):(\d+).*",line)
            try:
                ret.append(rexpl.groups())
            except AttributeError:
                pass
        return ret
    
    def flush(self):
        """ 
        Flushes the output buffer
        """
        if self.connected != True:
            raise ConnectionError
        try:
            for i in range(10):
                self.pipe.sendline("")
                self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
            raise TIMEOUT
        
    def exploit(self, include=None, exclude=None):
        """
        Starts exploiting of every service and every host existing
        in database. Optionally you can specify exclusion or/and inclusion,
        by passing particular network ranges as the "include" or "exclude" parameter.
        Example:
        MetaSploit.exploit(include="172.30.31.0/24")
         or 
        MetaSploit.exploit(include="10.0.0.1-10.0.0.50")
        
        "exclude" parameter syntax is same as for "include" one.
        """
        if self.connected != True:
            raise ConnectionError

        cmd = "db_autopwn -p -e"
        # exclusion and inclusion ranges
        if include:
            cmd += "-I %s" % include
        if exclude:
            cmd += "-E %s" % exclude 
        
        try:
                self.pipe.sendline(cmd)
                self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
                raise TIMEOUT
        lines = self.pipe.before.split("\r\n")
        
        # small hack to gather data from exploits
        # which has returned they results after msf> command prompt
        # has been displayed
        
        time.sleep(10)
        try:
                self.pipe.sendline("\n")
                self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
                raise TIMEOUT
        lines = lines + self.pipe.before.split("\r\n")
        return lines
    def listsessions(self):
        """Lists currently connected telnet sessions to the exploited hosts.
           Returns list of string containing raw output from MetaSploit."""
        if self.connected != True:
            raise ConnectionError
        try:
                self.pipe.expect("msf \>")
                self.pipe.sendline("sessions -l")
                self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
                raise TIMEOUT
        lines = self.pipe.before.split("\r\n")
        return lines
    def scan(self, network):
        """
        Runs scan for services on specified network. This command spawns
        nmap port scanner through MetaSploit framework. Notice that, if you
        want to run scan on remote framework, nmap installation on remote 
        host (running MSF) is required.
        """
        if self.connected != True:
            raise ConnectionError
        try:
                self.pipe.sendline("db_nmap %s" % network)
                self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
                raise TIMEOUT
        lines = self.pipe.before.split("\r\n")
        return lines
        
    def loadPlugins(self, login,password,dbname,hostname,port):
        """
        Loads all required framework modules. Currently used to load PostgreSQL backend
        module, then make connection to database using credentials passed to the connect() method.
        """
        self.pipe.expect("msf \>")
        self.pipe.sendline("load db_postgres")
        try:
            self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
            raise TIMEOUT
        if not re.search("Successfully loaded plugin", self.pipe.before):
            raise PluginError, "Problem loading db_postgres plugin for metasploit"
        self.pipe.sendline("db_connect %s:%s@%s:%s/%s" % (login,password,hostname,port,dbname))
        try:
            self.pipe.expect("msf \>")
        except pexpect.TIMEOUT:
            raise TIMEOUT
class MetaSploitLocal(MetaSploit):
    """
    Class containing code responsible for communication with locally spawned
    MetaSploit framework
    """
    def __init__(self,msfconsolepath):
        self.path = msfconsolepath
        self.port = None
        self.host = None
        self.connected = False
    def connect(self,login,password,dbname,hostname="localhost",port="5432"):
        try:
            self.pipe = pexpect.spawn(self.path, timeout=MSFTIMEOUT)
        except ExceptionPexpect:
            raise MsfExecNotFound
            
        self.connected = True
        self.loadPlugins(login,password,dbname,hostname,port)

class MetaSploitRemote(MetaSploit):
    """
    Class containing code responsible for communication with MetaSploit framework
    running on the remote host. Communication is done via TCP/IP->telnet protocol.
    """
    def __init__(self,host,port):
        self.host = host
        self.port = port
        self.path = None
        self.connected = False
    def connect(self,login,password,dbname,hostname="localhost",port="5432"):
        self.pipe = pexpect.spawn("telnet %s %s\n" % (self.host, self.port), timeout=MSFTIMEOUT)
        self.connected = True
        self.loadPlugins(login,password,dbname,hostname,port)